package com.k2fsa.sherpa.ncnn;

import java.util.Properties;

public class ModelConfig {
    private String encoderParam, encoderBin, decoderParam, decoderBin, joinerParam,
            joinerBin, tokens;
    private int numThreads = 1;
    private boolean useGPU = true;

    private static final int DEFAULT_THREAD_NUMBERS = 3;

    public static final String MODEL_DIR_KEY = "sherpa.ncnn.model.dir";
    private static final String MODEL_DIR;
    static {
        Properties props = System.getProperties ();
        if (props.containsKey (MODEL_DIR_KEY)) {
            MODEL_DIR = props.getProperty (MODEL_DIR_KEY);
        } else {
            MODEL_DIR = "../models";
        }
    }

    public ModelConfig () {
    }

    public ModelConfig (String encoderParam, String encoderBin, String decoderParam, String decoderBin,
                        String joinerParam, String joinerBin, String tokens, int numThreads, boolean useGPU) {
        this.encoderParam = encoderParam;
        this.encoderBin = encoderBin;
        this.decoderParam = decoderParam;
        this.decoderBin = decoderBin;
        this.joinerParam = joinerParam;
        this.joinerBin = joinerBin;
        this.tokens = tokens;
        this.numThreads = numThreads;
        this.useGPU = useGPU;
    }

    public String getEncoderParam () {
        return encoderParam;
    }

    public void setEncoderParam (String encoderParam) {
        this.encoderParam = encoderParam;
    }

    public String getEncoderBin () {
        return encoderBin;
    }

    public void setEncoderBin (String encoderBin) {
        this.encoderBin = encoderBin;
    }

    public String getDecoderParam () {
        return decoderParam;
    }

    public void setDecoderParam (String decoderParam) {
        this.decoderParam = decoderParam;
    }

    public String getDecoderBin () {
        return decoderBin;
    }

    public void setDecoderBin (String decoderBin) {
        this.decoderBin = decoderBin;
    }

    public String getJoinerParam () {
        return joinerParam;
    }

    public void setJoinerParam (String joinerParam) {
        this.joinerParam = joinerParam;
    }

    public String getJoinerBin () {
        return joinerBin;
    }

    public void setJoinerBin (String joinerBin) {
        this.joinerBin = joinerBin;
    }

    public String getTokens () {
        return tokens;
    }

    public void setTokens (String tokens) {
        this.tokens = tokens;
    }

    public int getNumThreads () {
        return numThreads;
    }

    public void setNumThreads (int numThreads) {
        this.numThreads = numThreads;
    }

    public boolean isUseGPU () {
        return useGPU;
    }

    public void setUseGPU (boolean useGPU) {
        this.useGPU = useGPU;
    }

    public static ModelConfig getModelConfig (int type, boolean useGPU) {
        switch (type) {
            case 0 : return createDefaultModelConfig (useGPU);
            case 1 : return createTransducerModelConfig (useGPU);
            case 2 : return createBilingualZhEnModelConfig (useGPU);
            case 3 : return createStreamingZipFormerEnModelConfig (useGPU);
            case 4 : return createStreamingZipFormerFrModelConfig (useGPU);
            case 5 : return createStreamingZipFormerBilingualZhEnModelConfig (useGPU);
            case 6 : return createStreamingZipFormerSmallBilingualZhEnModelConfig (useGPU);

            default: return null;
        }
    }

    // type = 0
    private static ModelConfig createDefaultModelConfig (boolean useGPU) {
        ModelConfig config = new ModelConfig ();
        String modelDir = MODEL_DIR + "/sherpa-ncnn-2022-09-30";
        config.encoderParam = modelDir + "/encoder_jit_trace-pnnx.ncnn.param";
        config.encoderBin = modelDir + "/encoder_jit_trace-pnnx.ncnn.bin";
        config.decoderParam = modelDir + "/decoder_jit_trace-pnnx.ncnn.param";
        config.decoderBin = modelDir + "/decoder_jit_trace-pnnx.ncnn.bin";
        config.joinerParam = modelDir + "/joiner_jit_trace-pnnx.ncnn.param";
        config.joinerBin = modelDir + "/joiner_jit_trace-pnnx.ncnn.bin";
        config.tokens = modelDir + "/tokens.txt";
        config.numThreads = DEFAULT_THREAD_NUMBERS;
        config.useGPU = useGPU;

        return config;
    }

    // type = 1
    private static ModelConfig createTransducerModelConfig (boolean useGPU) {
        String modelDir = MODEL_DIR + "/sherpa-ncnn-conv-emformer-transducer-2022-12-06";
        ModelConfig config = new ModelConfig ();
        config.encoderParam = modelDir + "/encoder_jit_trace-pnnx.ncnn.int8.param";
        config.encoderBin = modelDir + "/encoder_jit_trace-pnnx.ncnn.int8.bin";
        config.decoderParam = modelDir + "/decoder_jit_trace-pnnx.ncnn.param";
        config.decoderBin = modelDir + "/decoder_jit_trace-pnnx.ncnn.bin";
        config.joinerParam = modelDir + "/joiner_jit_trace-pnnx.ncnn.int8.param";
        config.joinerBin = modelDir + "/joiner_jit_trace-pnnx.ncnn.int8.bin";
        config.tokens = modelDir + "/tokens.txt";
        config.numThreads = DEFAULT_THREAD_NUMBERS;
        config.useGPU = useGPU;

        return config;
    }

    // type = 2
    private static ModelConfig createBilingualZhEnModelConfig (boolean useGPU) {
        String modelDir = MODEL_DIR + "/sherpa-ncnn-streaming-zipformer-bilingual-zh-en-2023-02-13";
        ModelConfig config = new ModelConfig ();
        config.encoderParam = modelDir + "/encoder_jit_trace-pnnx.ncnn.param";
        config.encoderBin = modelDir + "/encoder_jit_trace-pnnx.ncnn.bin";
        config.decoderParam = modelDir + "/decoder_jit_trace-pnnx.ncnn.param";
        config.decoderBin = modelDir + "/decoder_jit_trace-pnnx.ncnn.bin";
        config.joinerParam = modelDir + "/joiner_jit_trace-pnnx.ncnn.param";
        config.joinerBin = modelDir + "/joiner_jit_trace-pnnx.ncnn.bin";
        config.tokens = modelDir + "/tokens.txt";
        config.numThreads = DEFAULT_THREAD_NUMBERS;
        config.useGPU = useGPU;

        return config;
    }

    // type =3
    private static ModelConfig createStreamingZipFormerEnModelConfig (boolean useGPU) {
        String modelDir = MODEL_DIR + "/sherpa-ncnn-streaming-zipformer-en-2023-02-13";
        ModelConfig config = new ModelConfig ();
        config.encoderParam = modelDir + "/encoder_jit_trace-pnnx.ncnn.param";
        config.encoderBin = modelDir + "/encoder_jit_trace-pnnx.ncnn.bin";
        config.decoderParam = modelDir + "/decoder_jit_trace-pnnx.ncnn.param";
        config.decoderBin = modelDir + "/decoder_jit_trace-pnnx.ncnn.bin";
        config.joinerParam = modelDir + "/joiner_jit_trace-pnnx.ncnn.param";
        config.joinerBin = modelDir + "/joiner_jit_trace-pnnx.ncnn.bin";
        config.tokens = modelDir + "/tokens.txt";
        config.numThreads = DEFAULT_THREAD_NUMBERS;
        config.useGPU = useGPU;

        return config;
    }

    // type = 4
    private static ModelConfig createStreamingZipFormerFrModelConfig (boolean useGPU) {
        String modelDir = MODEL_DIR + "/sherpa-ncnn-streaming-zipformer-fr-2023-04-14";
        ModelConfig config = new ModelConfig ();
        config.encoderParam = modelDir + "/encoder_jit_trace-pnnx.ncnn.param";
        config.encoderBin = modelDir + "/encoder_jit_trace-pnnx.ncnn.bin";
        config.decoderParam = modelDir + "/decoder_jit_trace-pnnx.ncnn.param";
        config.decoderBin = modelDir + "/decoder_jit_trace-pnnx.ncnn.bin";
        config.joinerParam = modelDir + "/joiner_jit_trace-pnnx.ncnn.param";
        config.joinerBin = modelDir + "/joiner_jit_trace-pnnx.ncnn.bin";
        config.tokens = modelDir + "/tokens.txt";
        config.numThreads = DEFAULT_THREAD_NUMBERS;
        config.useGPU = useGPU;
        return config;
    }

    // type = 5
    private static ModelConfig createStreamingZipFormerBilingualZhEnModelConfig (boolean useGPU) {
        String modelDir = MODEL_DIR + "/sherpa-ncnn-streaming-zipformer-bilingual-zh-en-2023-02-13";
        ModelConfig config = new ModelConfig ();
        config.encoderParam = modelDir + "/encoder_jit_trace-pnnx.ncnn.param";
        config.encoderBin = modelDir + "/encoder_jit_trace-pnnx.ncnn.bin";
        config.decoderParam = modelDir + "/decoder_jit_trace-pnnx.ncnn.param";
        config.decoderBin = modelDir + "/decoder_jit_trace-pnnx.ncnn.bin";
        config.joinerParam = modelDir + "/joiner_jit_trace-pnnx.ncnn.param";
        config.joinerBin = modelDir + "/joiner_jit_trace-pnnx.ncnn.bin";
        config.tokens = modelDir + "/tokens.txt";
        config.numThreads = DEFAULT_THREAD_NUMBERS;
        config.useGPU = useGPU;

        return config;
    }

    // type = 6
    private static ModelConfig createStreamingZipFormerSmallBilingualZhEnModelConfig (boolean useGPU) {
        String modelDir = MODEL_DIR + "/sherpa-ncnn-streaming-zipformer-small-bilingual-zh-en-2023-02-16";
        ModelConfig config = new ModelConfig ();
        config.encoderParam = modelDir + "/encoder_jit_trace-pnnx.ncnn.param";
        config.encoderBin = modelDir + "/encoder_jit_trace-pnnx.ncnn.bin";
        config.decoderParam = modelDir + "/decoder_jit_trace-pnnx.ncnn.param";
        config.decoderBin = modelDir + "/decoder_jit_trace-pnnx.ncnn.bin";
        config.joinerParam = modelDir + "/joiner_jit_trace-pnnx.ncnn.param";
        config.joinerBin = modelDir + "/joiner_jit_trace-pnnx.ncnn.bin";
        config.tokens = modelDir + "/tokens.txt";
        config.numThreads = DEFAULT_THREAD_NUMBERS;
        config.useGPU = useGPU;

        return config;
    }
}